use std::fmt::Display;
use std::io::ErrorKind;

#[cfg(feature = "ser")]
use serde::de::DeserializeOwned;
#[cfg(feature = "ser")]
use serde::Serialize;

use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::{TcpListener, TcpSocket, TcpStream};

use crate::export::{
    BroadcastReceiver,
    BroadcastSender,
    MpscReceiver,
    MpscSender,
    yield_now,
    select
};

use crate::encrypt::{encrypt, decrypt, WTASK_KEY};

use crate::error::{WError, WResult};

static DATA_MAX_LEN: usize = 4096;
pub static TOKEN_START: &str = "winter";
pub static TOKEN_ALIVE: &str = "alive";


pub type WTcpStream = (OwnedReadHalf, OwnedWriteHalf);



#[derive(Debug, PartialEq, Eq)]
pub enum TCPClientType {
    Master,
    Session,
}

impl Display for TCPClientType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl TCPClientType {
    pub fn seq(&self, other: &str) -> bool {
        self.to_string().eq(other)
    }
}



pub async fn stream_read_raw(stream: &TcpStream, key: &str) -> WResult<Vec<u8>> {
    let mut data_raw = vec![0; 8];
    let mut data_all = Vec::new();
    let mut data_size = None;
    loop {
        stream.readable().await?;
        // log::debug!("read wait raw {}", data_raw.len());
        match stream.try_read(&mut data_raw) {
            Ok(0) => {
                return Err(WError::NetError("Channel close".to_owned()));
            }
            Ok(n) => {
                // log::debug!("tcp read {} data", n);
                if data_size.is_none() {
                    let t = data_raw.try_into().unwrap();
                    let t = u64::from_be_bytes(t) as usize;
                    data_size = Some(t);
                    let m = if t > DATA_MAX_LEN {
                        DATA_MAX_LEN
                    } else {
                        t
                    };
                    data_raw = vec![0; m];
                    continue;
                };
                let data_size = data_size.unwrap();
                if n <= data_size {
                    data_all.extend(data_raw.iter().take(n));
                }
                let m = data_size - data_all.len();
                if m == 0 {
                    break;
                }
                let m = if m > DATA_MAX_LEN {
                    DATA_MAX_LEN
                } else {
                    m
                };
                data_raw.truncate(m);
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                continue;
            }
            Err(e) => {
                return Err(e.into());
            }
        }
    }
    let data_raw = decrypt(&data_all, key)?;
    // log::debug!("stream read raw {} {} {:?}", data_all.len(), key, String::from_utf8_lossy(&data_raw));
    Ok(data_raw)
}


pub async fn stream_write_raw(stream: &TcpStream, data: &[u8], key: &str) -> WResult<()> {
    // log::debug!("stream write raw0 {} {} {:?}", data.len(), key, String::from_utf8_lossy(data));
    let data = encrypt(data, key)?;
    // log::debug!("stream write raw {} {}", data.len(), key);
    let data_len = (data.len() as u64).to_be_bytes().to_vec();
    let mut data_final = [data_len, data].concat();
    loop {
        // log::debug!("write wait raw");
        stream.writable().await?;
        match stream.try_write(&data_final) {
            Ok(n) => {
                if n != data_final.len() {
                    data_final = data_final.drain(n..).collect();
                    continue;
                }
                break;
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                continue;
            }
            Err(e) => {
                return Err(e.into());
            }
        }
    }
    Ok(())
}



pub async fn stream_read_token(stream: &TcpStream, key: &str) -> WResult<String> {
    let token_raw = stream_read_raw(stream, key).await?;
    let token = String::from_utf8_lossy(&token_raw).to_string();
    if !token.starts_with(WTASK_KEY) {
        return Err(WError::NetError("key error".to_owned()));
    }
    Ok(token[WTASK_KEY.len()..].to_string())
}



pub async fn stream_write_token<T: Display + Clone>(stream: &TcpStream, token: T, key: &str) -> WResult<()> {
    let data = format!("{WTASK_KEY}{token}");
    let data_raw = data.as_bytes();
    stream_write_raw(stream, data_raw, key).await
}





#[cfg(feature = "ser")]
pub async fn stream_write_data<V: Serialize + Send>(data: &V, stream_write: &TcpStream, key: &str) -> WResult<()> {
    let data = serde_json::to_vec(data)?;
    // log::debug!("sending data: {key}");
    stream_write_raw(stream_write, &data, key).await?;
    Ok(())
}

#[cfg(feature = "ser")]
pub async fn stream_read_data<V: DeserializeOwned + Send>(stream_read: &TcpStream, key: &str) -> WResult<V> {
    // log::debug!("reading data: {key}");
    let data = stream_read_raw(stream_read, key).await?;
    let res = serde_json::from_slice::<V>(&data)?;
    Ok(res)
}



pub fn tcp_listen(port: u32, tag: &str) -> WResult<TcpListener> {
    let ip = format!("0.0.0.0:{port}").parse()?;
    let socket = TcpSocket::new_v4()?;
    socket.set_reuseaddr(true)?;
    socket.bind(ip)?;
    let listener = socket.listen(1024)?;
    log::info!("Machine Listening on {ip}, {tag}");
    Ok(listener)
}





pub async fn tcp_client_start(ip: &str, port: u32, server_token: &str, tcp_type: &TCPClientType) -> WResult<WTcpStream> {
    let ip = format!("{ip}:{port}");
    // log::debug!("Connecting to server {}", &ip);
    let stream = TcpStream::connect(&ip).await?;
    // log::debug!("Connected to server {}", &ip);
    // log::debug!("Sending base token: {} {}", &ip, server_token);
    let stream_split = stream.into_split();
    stream_write_token(stream_split.1.as_ref(), server_token, WTASK_KEY).await?;
    // log::debug!("reading tcp type: {}", &ip);
    let tcp = stream_read_token(stream_split.0.as_ref(),server_token).await?;
    if tcp_type.seq(&tcp) {
        Ok(stream_split)
    } else {
        Err(WError::NetError("tcp type error".to_owned()))
    }
}



pub async fn tcp_client_session_start(ip: &str, port: u32, server_token: &str, session_token: &str) -> WResult<WTcpStream> {
    // log::debug!("tcp_client_session_start: {ip}");
    let stream = tcp_client_start(
        ip,
        port,
        server_token,
        &TCPClientType::Session
    ).await?;
    // log::debug!("tcp_client_session_start sending session token: {ip}");
    stream_write_token(stream.1.as_ref(), session_token, server_token).await?;
    // log::debug!("tcp_client_session_start sending session start: {ip}");
    stream_write_token(stream.0.as_ref(), TOKEN_START, session_token).await?;
    // log::debug!("tcp_client_session_start reading session start: {ip}");
    let token_start_res = stream_read_token(stream.0.as_ref(), session_token).await?;
    if !token_start_res.eq(TOKEN_START) {
        return Err(WError::DataError("Session start error".to_owned()));
    }
    Ok(stream)
}

pub async fn tcp_server_session_start(stream: &WTcpStream, server_token: &str) -> WResult<String> {
    let server_token_req = stream_read_token(stream.0.as_ref(), WTASK_KEY).await?;
    // log::debug!("local get server token {server_token_req:?}, {server_token:?}");
    if !server_token_req.eq(&server_token) {
        // log::error!("auth error: {server_token_req}");
        return Err(WError::DataError("Auth error".to_owned()));
    }
    stream_write_token(stream.1.as_ref(), &TCPClientType::Session.to_string(), server_token).await?;
    let session_token = stream_read_token(stream.0.as_ref(), server_token).await?;
    stream_write_token(stream.1.as_ref(), TOKEN_START, &session_token).await?;
    let token_start_res = stream_read_token(stream.0.as_ref(), &session_token).await?;
    if !token_start_res.eq(TOKEN_START) {
        return Err(WError::DataError("Session start error".to_owned()));
    }
    Ok(session_token)
}




async fn stream_exchange_dkey_inner(stream_read: &TcpStream, stream_write: &TcpStream, key: &str) -> WResult<()> {
    loop {
        if let Ok(data) = stream_read_raw(stream_read, key).await {
            let t = stream_write_raw(stream_write, &data, key).await;
            if let Err(e) = t {
                log::debug!("session stop: {e}");
                break;
            }
        } else {
            // log::debug!("session stop");
            break;
        }
        yield_now().await;
    }
    Ok(())
}

pub async fn stream_exchange_dkey(stream_a: &WTcpStream, stream_b: &WTcpStream, key: &str) -> WResult<()> {
    let (stream_a_read, stream_a_write) = stream_a;
    let (stream_b_read, stream_b_write) = stream_b;
    select! {
        _ = stream_exchange_dkey_inner(
            stream_b_read.as_ref(),
            stream_a_write.as_ref(),
            key
        ) => {},
        _ = stream_exchange_dkey_inner(
            stream_a_read.as_ref(),
            stream_b_write.as_ref(),
            key
        ) => {},
    }
    Ok(())
}




async fn stream_exchange_wkey_inner(stream_read: &TcpStream, stream_write: &TcpStream, key: &str, capacity: usize) -> WResult<()> {
    let mut data_raw = vec![0; capacity];
    loop {
        stream_read.readable().await?;
        match stream_read.try_read(&mut data_raw) {
            Ok(0) => {
                return Err(WError::NetError("Channel close".to_owned()));
            }
            Ok(n) => {
                let data = &data_raw[..n];
                let t = stream_write_raw(stream_write, data, key).await;
                if let Err(e) = t {
                    log::debug!("session stop: {e}");
                    break;
                }
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                continue;
            }
            Err(e) => {
                return Err(e.into());
            }
        }
        yield_now().await;
    }
    Ok(())
}

async fn stream_exchange_rkey_inner(stream_read: &TcpStream, stream_write: &TcpStream, key: &str) -> WResult<()> {
    loop {
        if let Ok(mut data) = stream_read_raw(stream_read, key).await {
            loop {
                stream_write.writable().await?;
                match stream_write.try_write(&data) {
                    Ok(n) => {
                        if n != data.len() {
                            data = data.drain(n..).collect();
                            continue;
                        }
                        break;
                    }
                    Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                        continue;
                    }
                    Err(e) => {
                        log::debug!("session error: {e}");
                        return Err(e.into());
                    }
                }
            }
        } else {
            log::debug!("session stop");
            break;
        }
        yield_now().await;
    }
    Ok(())
}

pub async fn stream_exchange_akey(stream_key: &WTcpStream, stream_normal: &WTcpStream, key: &str, capacity: usize) -> WResult<()> {
    let (stream_key_read, stream_key_write) = stream_key;
    let (stream_normal_read, stream_normal_write) = stream_normal;
    select! {
        _ = stream_exchange_wkey_inner(
            stream_normal_read.as_ref(),
            stream_key_write.as_ref(),
            key,
            capacity
        ) => {},
        _ = stream_exchange_rkey_inner(
            stream_key_read.as_ref(),
            stream_normal_write.as_ref(),
            key
        ) => {},
    }
    Ok(())
}



async fn stream_broadchannel_exchange_skey_rev_inner(stream_write: &TcpStream, mut receiver: BroadcastReceiver<Vec<u8>>, key: &str) -> WResult<()> {
    loop {
        let result = receiver.recv().await;
        if let Err(e) = result {
            log::debug!("broadchannel Receive error: {e}");
            break;
        }
        let result = result.unwrap();
        let t = stream_write_raw(stream_write, &result, key).await;
        if let Err(e) = t {
            log::debug!("broadchannel write errir, session stop, {e:?}");
            break;
        }
        yield_now().await;
    }
    Ok(())
}

async fn stream_broadchannel_exchange_skey_read_inner(stream_read: &TcpStream, sender: BroadcastSender<Vec<u8>>, key: &str) -> WResult<()> {
    loop {
        let v = stream_read_raw(stream_read, key).await;
        if let Ok(data) = v {
            let result = sender.send(data);
            if let Err(e) = result {
                log::debug!("broadchannel send error: {e}");
                break;
            }
            yield_now().await;
        } else {
            log::debug!("broadchannel read error, session stop: {:?}", v.err().unwrap());
            break;
        }
    }
    Ok(())
}

pub async fn stream_broadchannel_exchange_skey(stream: &WTcpStream, sender: BroadcastSender<Vec<u8>>, receiver: BroadcastReceiver<Vec<u8>>, key: &str) -> WResult<()> {
    let (stream_read, stream_write) = stream;
    select! {
        _ = stream_broadchannel_exchange_skey_rev_inner(
            stream_read.as_ref(),
            receiver,
            key
        ) => {},
        _ = stream_broadchannel_exchange_skey_read_inner(
            stream_write.as_ref(),
            sender,
            key
        ) => {},
    }
    Ok(())
}


async fn stream_mpscchannel_exchange_nokey_sread_inner(stream_read: &TcpStream, sender: MpscSender<Vec<u8>>, capacity: usize) -> WResult<()> {
    let mut data_raw = vec![0; capacity];
    loop {
        stream_read.readable().await?;
        match stream_read.try_read(&mut data_raw) {
            Ok(0) => {
                log::debug!("Channel close");
                break;
            }
            Ok(n) => {
                let data = &data_raw[..n];
                let result = sender.send(data.to_vec()).await;
                if let Err(e) = result {
                    log::debug!("session stop: {e}");
                    break;
                }
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                continue;
            }
            Err(e) => {
                log::debug!("session stop: {e}");
                break;
            }
        }
        yield_now().await;
    }
    Ok(())
}

async fn stream_mpscchannel_exchange_nokey_swrite_inner(stream_write: &TcpStream, mut receiver: MpscReceiver<Vec<u8>>) -> WResult<()> {
    loop {
        let data = receiver.recv().await;
        if data.is_none() {
            log::debug!("mpscchannel receive error, closed");
            break;
        }
        let mut data = data.unwrap();
        loop {
            stream_write.writable().await?;
            match stream_write.try_write(&data) {
                Ok(n) => {
                    if n != data.len() {
                        data = data.drain(n..).collect();
                        continue;
                    }
                    break;
                }
                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                    continue;
                }
                Err(e) => {
                    log::debug!("session error: {e}");
                    return Err(e.into());
                }
            }
        }
        yield_now().await;
    }
    Ok(())
}

pub async fn stream_mpscchannel_exchange_nokey(stream: &WTcpStream, sender: MpscSender<Vec<u8>>, receiver: MpscReceiver<Vec<u8>>, capacity: usize) -> WResult<()> {
    let (stream_read, stream_write) = stream;
    select! {
        _ = stream_mpscchannel_exchange_nokey_sread_inner(
            stream_read.as_ref(),
            sender,
            capacity
        ) => {},
        _ = stream_mpscchannel_exchange_nokey_swrite_inner(
            stream_write.as_ref(),
            receiver
        ) => {},
    }
    Ok(())
}



async fn stream_mpscchannel_exchange_key_sread_inner(stream_read: &TcpStream, sender: MpscSender<Vec<u8>>, key: &str) -> WResult<()> {
    loop {
        if let Ok(data) = stream_read_raw(stream_read, key).await {
            let result = sender.send(data.to_vec()).await;
            if let Err(e) = result {
                log::debug!("session stop: {e}");
                break;
            }
        } else {
            log::debug!("session stop");
            break;
        }
        yield_now().await;
    }
    Ok(())
}

async fn stream_mpscchannel_exchange_key_swrite_inner(stream_write: &TcpStream, mut receiver: MpscReceiver<Vec<u8>>, key: &str) -> WResult<()> {
    loop {
        let data = receiver.recv().await;
        if data.is_none() {
            log::debug!("mpscchannel receive error, closed");
            break;
        }
        let data = data.unwrap();
        let result = stream_write_raw(stream_write, &data, key).await;
        if let Err(e) = result {
            log::debug!("session stop: {e}");
            break;
        }
        yield_now().await;
    }
    Ok(())
}

pub async fn stream_mpscchannel_exchange_key(stream: &WTcpStream, sender: MpscSender<Vec<u8>>, receiver: MpscReceiver<Vec<u8>>, key: &str) -> WResult<()> {
    let (stream_read, stream_write) = stream;
    select! {
        _ = stream_mpscchannel_exchange_key_sread_inner(
            stream_read.as_ref(),
            sender,
            key
        ) => {},
        _ = stream_mpscchannel_exchange_key_swrite_inner(
            stream_write.as_ref(),
            receiver,
            key
        ) => {},
    }
    Ok(())
}
